// pstree.go
package pstree

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

type Process struct {
	PID      int
	Name     string
	PPID     int
	CmdLine  string
	Children []int
}

type ProcessTree struct {
	Procs map[int]*Process
}

func BuildTree() (*ProcessTree, error) {
	procs := make(map[int]*Process)
	var errs []error

	files, err := filepath.Glob("/proc/[0-9]*")
	if err != nil {
		return nil, fmt.Errorf("扫描/proc失败: %v", err)
	}

	for _, dir := range files {
		pid := parsePID(dir)
		if pid == 0 {
			continue
		}

		proc := &Process{PID: pid}
		
		// 解析进程基本信息
		if statPath := filepath.Join(dir, "stat"); parseStat(statPath, proc) != nil {
			errs = append(errs, fmt.Errorf("解析stat失败: %s", statPath))
			continue
		}
		
		// 解析命令行参数
		if cmdlinePath := filepath.Join(dir, "cmdline"); parseCmdline(cmdlinePath, proc) != nil {
			errs = append(errs, fmt.Errorf("解析cmdline失败: %s", cmdlinePath))
		}

		procs[pid] = proc
	}

	// 构建进程树结构
	tree := &ProcessTree{Procs: procs}
	for _, proc := range procs {
		if parent, exists := procs[proc.PPID]; exists {
			parent.Children = append(parent.Children, proc.PID)
		}
	}
	
	return tree, errors.Join(errs...)
}

// 辅助函数 --------------------------------------------------
func parsePID(path string) int {
	base := filepath.Base(path)
	pid, err := strconv.Atoi(base)
	if err != nil || pid <= 0 {
		return 0
	}
	return pid
}

func parseStat(path string, proc *Process) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}

	// 示例数据：12345 (bash) S 6789 ...
	raw := string(data)
	commEnd := strings.IndexRune(raw, ')')
	if commEnd == -1 {
		return errors.New("无效stat格式")
	}

	// 提取进程名
	proc.Name = strings.Trim(raw[:commEnd], "(")

	// 解析PPID
	fields := strings.Fields(raw[commEnd+2:])
	if len(fields) < 3 {
		return errors.New("stat字段不足")
	}
	
	ppid, err := strconv.Atoi(fields[1])
	if err != nil {
		return fmt.Errorf("解析PPID失败: %v", err)
	}
	proc.PPID = ppid
	
	return nil
}

func parseCmdline(path string, proc *Process) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	
	// 处理空命令行的情况
	if len(data) == 0 {
		proc.CmdLine = ""
		return nil
	}
	
	// 替换null字符为空格
	proc.CmdLine = strings.ReplaceAll(string(data), "\x00", " ")
	return nil
}

func (tree *ProcessTree) FormatChain(pid int) string {
	var chain []string
	for pid != 0 {
		proc, exists := tree.Procs[pid]
		if !exists {
			break
		}
		chain = append([]string{fmt.Sprintf("%d(%s)", proc.PID, proc.Name)}, chain...)
		pid = proc.PPID
	}
	return strings.Join(chain, " -> ")
}

